#! /usr/bin/python3
# -*- coding: utf-8 -*-
"""
Created on 2020/1/14 8:49 星期二

@author: jyz
"""
import os
import time
import signal
import requests

import tensorflow as tf
from tensorflow.compat.v1 import ConfigProto, GraphDef
from tensorflow.compat.v1 import InteractiveSession
import cv2
import numpy as np

from config.config_setting import MODEL_PATH, DEVICE_PIDS, STREAM_SIZE
from config.config_setting import logger
from model import crowded_detection, cross_line_detection, fire_detection, ref_clothes_detection, safety_helmet_detection, bare_soil_detection, fall_down_detection
from model import model_on_gpu
from model import model_on_gpu_trt
"""该文件内主要存放启动各模型需要用到的方法"""

def change_device_by_model(origin_img_q, result_img_q, algo_id, config_data):
    if algo_id in [10, 20, 30, 40, 80, 90]:
        model_on_gpu_trt.prepare_trt(origin_img_q, result_img_q, algo_id, config_data)
    elif algo_id == 50:
        model_on_gpu.fall_model(origin_img_q, result_img_q, algo_id, config_data)
    else:
        model_on_gpu.tf_model(origin_img_q, result_img_q, algo_id, config_data)
        # tf_model_2(origin_img_q, result_img_q, algo_id, config_data)


# 模型加载和推理
# 算法id和功能对应列表{安全帽：10，反光衣：20，越界：30，人员聚集：40，摔倒行为检测：50，人员离岗：60，人员徘徊：70，裸土覆盖：80，火焰检测：90}
def tf_model(origin_img_q, result_img_q, algo_id, reqs):
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

    tf_config = ConfigProto(allow_soft_placement=True)
    time_ = time.time()

    # 加载模型
    '''loading the TensorFlow model'''
    ckpt_pwd = os.path.join(MODEL_PATH[algo_id], 'frozen_inference_graph.pb')
    if not os.path.exists(ckpt_pwd):   # 当本地没有模型时，自动下载。
        model_url = "https://whfciot.oss-cn-hangzhou.aliyuncs.com/ai_camera/"
        model_req = requests.get(os.path.join(model_url, ckpt_pwd))
        with open(ckpt_pwd, 'wb') as cod:
            cod.write(model_req.content)

    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = GraphDef()
        with tf.gfile.GFile(ckpt_pwd, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    '''run the TensorFlow model'''
    with detection_graph.as_default():
        '''limit the GPU Memory'''
        with tf.Session(graph=detection_graph, config=tf_config) as sess:
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
            logger.info("||| Loading TF model time:{}".format(time.time() - time_))

            while True:
                for devid in reqs:
                    if origin_img_q[devid].qsize() == 0:
                        time.sleep(0.1)
                        continue
                    image = origin_img_q[devid].get()

                    # 判断在当前时间算法是否可运行
                    detect_result = None
                    timer = reqs[devid][algo_id]["timeList"][0] if "timeList" in reqs[devid][algo_id].keys() else []
                    if timer:
                        timer_s, timer_e = timer[0]["startTime"], timer[0]["endTime"]  # 提取起止时间
                        if timer_s <= time.strftime('%H:%M') <= timer_e:
                            if image is not None:
                                if min(image.shape[:2]) < 100:
                                    logger.error("May the image of camera is None.")
                                    continue
                                img_ori = image.copy()
                                img_ori = cv2.cvtColor(img_ori, cv2.COLOR_BGR2RGB)

                                detect_result = sess.run(
                                    [detection_boxes, detection_scores, detection_classes, num_detections],
                                    feed_dict={image_tensor: img_ori[np.newaxis, :, :, :]},
                                )
                    else:
                        if image is not None:
                            if min(image.shape[:2]) < 100:
                                logger.error("May the image of camera is None.")
                                continue
                            img_ori = image.copy()
                            img_ori = cv2.cvtColor(img_ori, cv2.COLOR_BGR2RGB)
                            try:  # 双重保险，防止image的尺寸过小导致报错
                                # tic = datetime.now()
                                detect_result = sess.run(
                                    [detection_boxes, detection_scores, detection_classes, num_detections],
                                    feed_dict={image_tensor: img_ori[np.newaxis, :, :, :]},)  # (boxes, scores, label_id, nums)
                                # toc = datetime.now()
                                # print('Elapsed time: %f seconds' % (toc-tic).total_seconds())
                            except TypeError:
                                pass

                    if detect_result is not None:
                        result_img_q[devid].put((image, detect_result))
                    else:
                        result_img_q[devid].put(image)


def select_algo_draw_box(boxes, scores, label_id, label_dict, img_cur, param, device_id, algo):
    if algo == 10:    #安全帽
        return safety_helmet_detection.draw_box_by_data_trt(boxes, scores, label_id, label_dict, img_cur, device_id)
    elif algo == 20:  #反光衣
        return ref_clothes_detection.draw_box_by_data_trt(boxes, scores, label_id, label_dict, img_cur, device_id)
    elif algo == 30:  #人员越界
        return cross_line_detection.draw_box_by_data_trt(boxes, scores, label_id, label_dict, img_cur, param, device_id)
    elif algo == 40:  #人员聚集
        return crowded_detection.draw_box_by_data_trt(boxes, scores, label_id, label_dict, img_cur, param, device_id)
    elif algo == 50:  # 摔倒行为检测
        pass
    elif algo == 60:  # 人员离岗
        pass
    elif algo == 70:  # 人员徘徊检测
        pass
    elif algo == 80:  # 裸土覆盖检测
        return bare_soil_detection.draw_box_by_data_trt(boxes, scores, label_id, label_dict, img_cur, device_id)
    elif algo == 90:  # 火焰检测
        return fire_detection.draw_box_by_data_trt(boxes, scores, label_id, label_dict, img_cur, device_id)
    else:  # 车辆检测
        pass


def filter_algos(algos, rtmp_paths):
    algo_new, rtmp_new = [], []

    for algo in list(algos.keys()):
        if algo in [10, 20, 30, 40, 50, 80, 90]:
            algo_new.append(algo)
            rtmp_new.append(rtmp_paths[algo])
        else:
            logger.error("the algorithm {} haven't development for use", algo)
    return len(algo_new), algo_new, rtmp_new


# 根据算法选择检测结果判断方式
def queue_img_get(q_get, q_put, algo_id, device_config):
    if algo_id == 50:
        fall_down_detection.queue_img_get(q_get, q_put, device_config)
    elif algo_id in [10, 20, 30, 40, 80, 90]:
        comm_img_get_trt(q_get, q_put, algo_id, device_config)
    else:
        comm_img_get(q_get, q_put, algo_id, device_config)


def comm_img_get(q_get, q_put, algo_id, reqs):
    pbtxt_name = [f for f in os.listdir(MODEL_PATH[algo_id]) if f.find('.pbtxt') >= 0][0]
    with open(os.path.join(MODEL_PATH[algo_id], pbtxt_name), 'r') as f:  # get label_dict
        lines = f.readlines()
        ids = [int(line.split(':')[-1]) for line in lines if line.find('id:') > 0]
        names = [line.split(':')[-1][2:-2] for line in lines if line.find('display_name:') > 0]
        label_dict = dict(zip(ids, names))

    while True:
        for devid in reqs:
            if q_get[devid].qsize() == 0:
                time.sleep(0.1)
                continue
            
            img_info = q_get[devid].get()

            # 判断在当前时间算法是否可运行
            detect_result = None
            timer = reqs[devid][algo_id]["timeList"][0] if "timeList" in reqs[devid][algo_id].keys() else []
            if timer:
                timer_s, timer_e = timer[0]["startTime"], timer[0]["endTime"]  # 提取起止时间
                if timer_s <= time.strftime('%H:%M') <= timer_e:
                    if isinstance(img_info, tuple):
                        img, (boxes, scores, label_id, nums) = img_info
                        detect_result = select_algo_draw_box(boxes, scores, label_id, label_dict, img, reqs[devid][algo_id], devid, algo_id)
            else:
                if isinstance(img_info, tuple):
                    img, (boxes, scores, label_id, nums) = img_info
                    detect_result = select_algo_draw_box(boxes, scores, label_id, label_dict, img, reqs[devid][algo_id], devid, algo_id)
                    # cv2.imwrite("res.png", detect_result)
            if detect_result is not None:
                q_put[devid].put(detect_result)
            else:  # 如果不在可运行的时间范围内，则返回原图片帧
                q_put[devid].put(img_info)


def comm_img_get_trt(q_get, q_put, algo_id, reqs):
    pbtxt_name = [f for f in os.listdir(MODEL_PATH[algo_id]) if f.find('.pbtxt') >= 0][0]
    with open(os.path.join(MODEL_PATH[algo_id], pbtxt_name), 'r') as f:  # get label_dict
        lines = f.readlines()
        ids = [int(line.split(':')[-1]) for line in lines if line.find('id:') > 0]
        names = [line.split(':')[-1][2:-2] for line in lines if line.find('display_name:') > 0]
        label_dict = dict(zip(ids, names))

    while True:
        for devid in reqs:
            if q_get[devid].qsize() == 0:
                time.sleep(0.1)
                continue
            
            img_info = q_get[devid].get()

            # 判断在当前时间算法是否可运行
            detect_result = None
            timer = reqs[devid][algo_id]["timeList"][0] if "timeList" in reqs[devid][algo_id].keys() else []
            if timer:
                timer_s, timer_e = timer[0]["startTime"], timer[0]["endTime"]  # 提取起止时间
                if timer_s <= time.strftime('%H:%M') <= timer_e:
                    if isinstance(img_info, tuple):
                        img, (boxes, scores, label_id) = img_info
                        detect_result = select_algo_draw_box(boxes, scores, label_id, label_dict, img, reqs[devid][algo_id], devid, algo_id)
            else:
                if isinstance(img_info, tuple):
                    img, (boxes, scores, label_id) = img_info
                    detect_result = select_algo_draw_box(boxes, scores, label_id, label_dict, img, reqs[devid][algo_id], devid, algo_id)
                    # cv2.imwrite("res.png", detect_result)
            if detect_result is not None:
                q_put[devid].put(detect_result)
            else:  # 如果不在可运行的时间范围内，则返回原图片帧
                q_put[devid].put(img_info)

# kill需要重启的设备已经启动的算法进程。
def kill_started_pro(device):
    for p in DEVICE_PIDS[device]:
        os.kill(p, signal.SIGKILL)
    print('========', DEVICE_PIDS[device])
    DEVICE_PIDS[device] = []

